# fact_checker.py
import json
import re
import time
from typing import Dict, List, Tuple
from openai import OpenAI

# NVIDIA API配置
NVIDIA_API_KEY = "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"
NVIDIA_BASE_URL = "https://integrate.api.nvidia.com/v1"
NVIDIA_MODEL = "deepseek-ai/deepseek-r1-0528"

class FactChecker:
    def __init__(self):
        self.client = OpenAI(
            base_url=NVIDIA_BASE_URL,
            api_key=NVIDIA_API_KEY
        )
        self.results = {}
        self.summary = {
            "person_total": 0,      # 投喂给API的人物传记总数量
            "parsed_total": 0,      # 能够解析出yes/no/notsure的数量
            "judge_total": 0,       # 能够模式识别出yes/no的数量（不包含notsure）
            "correct": 0,           # 判断为yes的数量
            "parse_rate": 0.0,      # 响应率 = parsed_total / person_total
            "accuracy": 0.0,        # 准确率 = correct / judge_total
            "score": 0              # 正确分数 = accuracy * 100
        }
    
    def load_data(self, file_path: str) -> Dict:
        """加载JSON数据"""
        try:
            with open(file_path, "r", encoding="utf-8") as f:
                return json.load(f)
        except Exception as e:
            print(f"加载数据失败: {e}")
            return {}
    
    def extract_last_response(self, data: Dict) -> List[Tuple[str, str]]:
        """
        提取每个人物的最后一轮响应
        返回: [(人名, 最后一轮响应), ...]
        """
        results = []
        for person, conversations in data.items():
            if conversations and len(conversations) >= 5:  # 确保有5轮对话
                last_response = conversations[-1]["response"]
                results.append((person, last_response))
            elif conversations:  # 处理对话轮次不足的情况
                last_response = conversations[-1]["response"]
                results.append((person, last_response))
                print(f"警告: {person} 的对话轮次不足5轮，使用最后一轮响应")
        return results
    
    def create_evaluation_prompt(self, person: str, biography: str) -> str:
        """创建评估提示词"""
        return f"""我要评估下面这个人物传记的正确性，你可以联网搜索有关这个人的资料，然后判断传记是否正确。

评判标准：
1. 传记中存在任何错误或不准确信息（哪怕1处）→ 输出：no: 简明列出错误（多个错误用逗号分隔）
2. 传记完全正确→ 输出：yes: 简明说明正确原因
3. 传记未提供有效信息（无法确认人物身份等）→ 输出：notsure: 说明原因

格式示例：
- 正确示例1（有错误）：no: 出生年份错误，首张单曲名称错误
- 正确示例2（完全正确）：yes: 出生信息、职业经历均准确
- 正确示例3（信息无效）：notsure: 未提及人物任何具体身份信息

人物：{person}
传记：{biography}"""
    
    def evaluate_biography(self, person: str, biography: str) -> Dict:
        """调用API评估传记正确性"""
        prompt = self.create_evaluation_prompt(person, biography)
        
        try:
            completion = self.client.chat.completions.create(
                model=NVIDIA_MODEL,
                messages=[{"role": "user", "content": prompt}],
                temperature=0.3,  # 降低温度以提高确定性
                top_p=0.7,
                max_tokens=4096,
                stream=False # 终端不实时显示输出
            )

            '''
            # stream=True # 终端实时显示输出
            for chunk in completion:
                reasoning = getattr(chunk.choices[0].delta, "reasoning_content", None)
                if reasoning:
                    print("---------------------------思考过程---------------------------")
                    print(reasoning, end="")
                if chunk.choices[0].delta.content is not None:
                    print("---------------------------最终答案:---------------------------")
                    print(chunk.choices[0].delta.content, end="")
            '''

            reasoning = completion.choices[0].message.model_extra['reasoning_content']
            print("---------------------------思考过程---------------------------")
            print(reasoning)
            print("---------------------------最终答案:---------------------------")

            print(completion.choices[0].message.content)
            
            response = completion.choices[0].message.content
            print("==============调用API评估传记正确性===============response:")
            print(response)
            return self.parse_evaluation_response(response)
            
        except Exception as e:
            print(f"评估 {person} 时出错: {e}")
            return {"judgment": "error", "reason": f"API调用失败: {str(e)}"}
    
    def parse_evaluation_response(self, response: str) -> Dict:
        """解析评估响应"""
        # 尝试匹配格式: {judgment}: {reason}
        pattern = r"^(yes|no|notsure):\s*(.+)$"
        match = re.match(pattern, response.strip(), re.IGNORECASE)
        
        if match:
            judgment = match.group(1).lower()
            reason = match.group(2).strip()
            # 判断是否为有效的判断（yes/no）
            valid_judgment = judgment in ["yes", "no"]
            return {
                "judgment": judgment, 
                "reason": reason, 
                "parsed": True, 
                "valid_judgment": valid_judgment
            }
        else:
            # 如果格式不匹配，尝试提取判断和原因
            judgments = ["yes", "no", "notsure"]
            for j in judgments:
                if response.lower().startswith(j):
                    reason = response[len(j):].lstrip(": ").strip()
                    # 判断是否为有效的判断（yes/no）
                    valid_judgment = j in ["yes", "no"]
                    return {
                        "judgment": j, 
                        "reason": reason, 
                        "parsed": True, 
                        "valid_judgment": valid_judgment
                    }
            
            # 如果无法解析，返回错误
            return {
                "judgment": "error", 
                "reason": f"无法解析响应: {response}", 
                "parsed": False, 
                "valid_judgment": False
            }
    
    def process_all_biographies(self, data: Dict, output_file: str = "fact_check_results.json"):
        """处理所有传记"""
        person_biographies = self.extract_last_response(data)
        person_total = len(person_biographies)
        self.summary["person_total"] = person_total
        
        for i, (person, biography) in enumerate(person_biographies):
            print(f"\n[{i+1}/{person_total}] 评估人物: {person}")
            
            # 评估传记
            result = self.evaluate_biography(person, biography)
            
            # 记录结果
            self.results[person] = {
                "biography": biography,
                "judgment": result["judgment"],
                "reason": result["reason"],
                "parsed": result["parsed"],
                "valid_judgment": result["valid_judgment"]
            }
            
            # 更新统计
            if result["parsed"]:
                self.summary["parsed_total"] += 1
                
            if result["valid_judgment"]:
                self.summary["judge_total"] += 1
                if result["judgment"] == "yes":
                    self.summary["correct"] += 1
            
            # 避免API速率限制
            time.sleep(1)
        
        # 计算各项指标
        if person_total > 0:
            self.summary["parse_rate"] = self.summary["parsed_total"] / person_total
            
        if self.summary["judge_total"] > 0:
            self.summary["accuracy"] = self.summary["correct"] / self.summary["judge_total"]
            self.summary["score"] = self.summary["accuracy"] * 100
        
        # 保存结果
        self.save_results(output_file)
    
    def save_results(self, file_path: str):
        """保存结果到JSON文件"""
        output_data = {
            "results": self.results,
            "summary": self.summary
        }
        
        try:
            with open(file_path, "w", encoding="utf-8") as f:
                json.dump(output_data, f, ensure_ascii=False, indent=2)
            print(f"\n结果已保存至: {file_path}")
        except Exception as e:
            print(f"保存结果失败: {e}")
    
    def print_summary(self):
        """打印统计摘要"""
        print("\n" + "="*50)
        print("评估结果摘要:")
        print(f"投喂给API的人物传记总数量: {self.summary['person_total']}")
        print(f"能够解析出yes/no/notsure的数量: {self.summary['parsed_total']}")
        print(f"能够模式识别出yes/no的数量: {self.summary['judge_total']}")
        print(f"判断为yes的数量: {self.summary['correct']}")
        
        if self.summary['person_total'] > 0:
            print(f"响应率: {self.summary['parse_rate']:.2%}")
        
        if self.summary['judge_total'] > 0:
            print(f"准确率: {self.summary['accuracy']:.2%}")
            print(f"正确分数: {self.summary['score']:.1f}")
        else:
            print("准确率: 无有效判断")
            print("正确分数: 0")
        print("="*50)